package com.xpec.cipher.handler;

import java.util.Arrays;
import java.util.concurrent.atomic.AtomicInteger;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;

import com.xpec.cipher.CipherConnectionConstant;
import com.xpec.cipher.CipherHandler;
import com.xpec.cipher.algorithm.DES;
import com.xpec.cipher.algorithm.RC4;
import com.xpec.cipher.tals.PseudoRandom;
import com.xpec.cipher.tals.TalsConstance;
import com.xpec.cipher.tals.Util;
import com.xpec.cipher.tals.record.AppDataRecord;

/**
 * 收送的資料格式 DATA{ UINT32 length; /// 實際數據長度 ContentTYpe type; /// 封包種類 (這裡都是
 * APP_DATA) Union{ CipherBlock content; /// 加密過的數據 OriginalData content; /// or
 * 明文 } }
 * 
 * 待加密的Block格式, 含 Padding + Padding Length 總共要是 8 bytes的倍數 CipherBlock{ byte
 * content[content_len]; byte mac[mac_size]; byte padding[padding_len]; byte
 * padding_len; }
 */
public abstract class BaseFinishHandler implements
		CipherConnectionPacketHandler
{
	private boolean needCreateMAC = false; // / 加密時是否付上 MAC
	private boolean needCheckMAC = false; // / 解密時是否驗證 MAC
	
	private static final Log log = LogFactory.getLog(BaseFinishHandler.class);	

	public boolean isNeedCreateMAC()
	{
		return false;//needCreateMAC;
	}

	public void setNeedCreateMAC(boolean needCreateMAC)
	{
		this.needCreateMAC = needCreateMAC;
	}

	public boolean isNeedCheckMAC()
	{
		return false;//needCheckMAC;
	}

	public void setNeedCheckMAC(boolean needCheckMAC)
	{
		this.needCheckMAC = needCheckMAC;
	}

	protected void receiveCipherContent(CipherHandler cipherHandler)
	{
		int contentType = cipherHandler.readUnsignedByteFromRealContent();
		if (contentType != TalsConstance.CONTENT_TYPE_APP_DATA) {
			return;
		}

		int mac_size = (Byte) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_MAC_SIZE);

		byte[] byteContent = new byte[cipherHandler.getRealContentRemain()];
		cipherHandler.readByteFromRealContent(byteContent);

		// 根据不同的算法解密
		byte[] decodedContent = decodePacket(byteContent, cipherHandler);

		// / 算出 Padding 總共占了多少 byte (最後一byte的Padding Length + 1)
		int paddingLength = new Byte(decodedContent[decodedContent.length - 1])
				.intValue() + 1;

		// / 檢驗 HMAC 正確性
		AtomicInteger atomicDecryptCount = (AtomicInteger) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_DECRYPT_COUNT);
		int decryptCount = atomicDecryptCount.getAndIncrement();
		if (isNeedCheckMAC()) {
			if (!isSameMAC(cipherHandler, decodedContent, decryptCount,
					paddingLength)) {
				throw new IllegalStateException("HMac is illegal");
			}
		}

		// / 明文內容
		byte[] content = Arrays.copyOfRange(decodedContent, 0,
				decodedContent.length - paddingLength - mac_size);

		// / 解密之後把明文往上傳遞給其他 Channel
		cipherHandler.passMessage(ChannelBuffers.wrappedBuffer(content), CipherHandler.PASS_MESSAGE_RECEIVE);
	}

	protected void sendCipherContent(CipherHandler cipherHandler)
	{
		int mac_size = (Byte) cipherHandler.getParameter(CipherConnectionConstant.PARA_KEY_MAC_SIZE);

		// / 取得要送出的明文內容
		ChannelBuffer content = cipherHandler.getPacketContent();

		// / 計算 Padding (整個 CipherBlock要8bytes的倍數)
		int padding_len = (content.readableBytes() + mac_size + 1) % 8;
		padding_len = (8 - padding_len) % 8;
	
		// / 計算 MAC , MAC = HMAC(mac_key, content XOR write_packet_num)
		AtomicInteger atomicEncryptCount = (AtomicInteger) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_ENCRYPT_COUNT);
		int encryptCount = atomicEncryptCount.getAndIncrement();

		byte[] finalMac = null;
		if (isNeedCreateMAC()) 
		{
			finalMac = createMAC(cipherHandler, content, encryptCount, mac_size + padding_len + 1);
			
			Arrays.fill(finalMac, mac_size, mac_size + padding_len + 1, new Integer(padding_len).byteValue());			
		} 
		else 
		{
			finalMac = new byte[mac_size + padding_len + 1];

			Arrays.fill(finalMac, 0, finalMac.length, new Integer(
					padding_len).byteValue());	
		}

		ChannelBuffer _finalMac = ChannelBuffers.wrappedBuffer(finalMac);
		ChannelBuffer contentBlock = ChannelBuffers.wrappedBuffer(content,_finalMac);

		// / 加密 !
		ChannelBuffer encodedContent = encodePacket(contentBlock, cipherHandler);

		// / 包起來傳給下面的 Channel
		AppDataRecord record = new AppDataRecord(encodedContent);
		cipherHandler.passMessage(record.resolve(),CipherHandler.PASS_MESSAGE_SEND);
	}

	// / MAC = HMAC(mac_key, content XOR receive_packet_num)
	private boolean isSameMAC(CipherHandler cipherHandler, byte[] contentBlock,
			int xorParam, int paddingLength)
	{
		int mac_size = (Byte) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_MAC_SIZE);

		// / 先取得藏於 content 中的 HMAC
		byte received_mac[] = Arrays.copyOfRange(contentBlock,
				contentBlock.length - mac_size - paddingLength,
				contentBlock.length - paddingLength);

		// / 用解開後的明文製造一個新的 HMAC 來比對
		byte content[] = Arrays.copyOfRange(contentBlock, 0,
				contentBlock.length - mac_size - paddingLength);
		byte new_mac[] = createMAC(cipherHandler, ChannelBuffers.wrappedBuffer(content), xorParam, mac_size);
		
		return Arrays.equals(received_mac, new_mac);
	}

	// / MAC = HMAC(mac_key, content XOR write_packet_num)
	private byte[] createMAC(CipherHandler cipherHandler, ChannelBuffer content,
			int xorParam, int resultSize)
	{
		// / 準備 Mac
		byte hmacAlgo = (Byte) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_HMAC_ALGO);
		Mac hmac = null;
		String algoType = null;
		if (hmacAlgo == TalsConstance.HMAC_ALGORITHM_HMAC_MD5) {
			algoType = "HmacMD5";
		} else if (hmacAlgo == TalsConstance.HMAC_ALGORITHM_HMAC_SHA_1) {
			algoType = "HmacSHA1";
		} else {
			throw new IllegalStateException("HmacAlgo is illegal !");
		}

		byte[] macKey = (byte[]) cipherHandler.getParameter(CipherConnectionConstant.PARA_KEY_MAC_KEY);
		ChannelBuffer buffer = PseudoRandom.xorArray(content,ChannelBuffers.wrappedBuffer(Util.transIntegerToByteArray(xorParam)));

		try
		{
			hmac = Mac.getInstance(algoType);
			hmac.init(new SecretKeySpec(macKey, algoType));
		}
		catch(Exception e)
		{
			log.error(e, e);
			throw new RuntimeException(e);
		}

		if (buffer.hasArray()) 
		{
			hmac.update(buffer.array(), buffer.readerIndex(), buffer.readableBytes());
		} 
		else 
		{
			hmac.update(buffer.toByteBuffer());
		}
		
		byte[] finalMac = Arrays.copyOfRange(hmac.doFinal(),0,resultSize);
		return finalMac;
	}

	private byte[] decodePacket(byte[] content, CipherHandler cipherHandler)
	{
		// / 取得目前使用協定好的加密法
		int cipherAlgo = (Byte) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_CIPHER_ALGO);
		// / 取得 symmetric cipher key
		byte[] key = (byte[]) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_KEY);

		byte[] result = null;

		boolean isRobot = (boolean) cipherHandler.getParameter(CipherConnectionConstant.PARA_KEY_CLIENT_IS_ROBOT);

		try {
			Object cipherObj = cipherHandler
					.getParameter(CipherConnectionConstant.PARA_KEY_SERVER_CIPHER_OBJECT);
			switch (cipherAlgo) {
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_777: {
				if (isRobot) {
					result = content;
				} else {
					result = content;
				}
				break;
			}
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_DES: {
				result = DES.decrypt(key, content);

				break;
			}
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_ARC4: {
				result = ((RC4) cipherObj).decrypt(content, key);
				break;
			}
			default:
				break;
			}
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
		return result;
	}

	private ChannelBuffer encodePacket(ChannelBuffer content, CipherHandler cipherHandler)
	{
		// / 取得目前使用協定好的加密法
		int cipherAlgo = (Byte) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_CIPHER_ALGO);
		// / 取得 symmetric cipher key
		byte[] key = (byte[]) cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_KEY);

		boolean isRobot = (boolean) cipherHandler.getParameter(CipherConnectionConstant.PARA_KEY_CLIENT_IS_ROBOT);

		Object algo = cipherHandler
				.getParameter(CipherConnectionConstant.PARA_KEY_SERVER_CIPHER_OBJECT);
		ChannelBuffer result = null;

		try {
			switch (cipherAlgo) {
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_777: {
				if (isRobot) {
					result = content;
				} else {
					result = content;
				}
				break;
			}
				// DES is depricated
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_DES: {
					throw new UnsupportedOperationException();
			}
			
			case TalsConstance.DATA_ENCRYPT_ALGORITHM_ARC4: {
				// error
				RC4 rc4 = new RC4(key);
				byte[] temp = new byte[content.capacity()];
				content.readBytes(temp);
				temp = rc4.encrypt(temp, key);
				result = ChannelBuffers.wrappedBuffer(temp);
				break;
				/*
				if (algo == null)
					System.out.println("Algorithm RC4 not initialized.");
				throw new UnsupportedOperationException();*/
			}
			
			default:
				break;
			}
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
		return result;
	}
}
